import orjson
from typing import Optional, Callable
from kafka import KafkaProducer, KafkaConsumer, KafkaAdminClient
from kafka.admin import NewTopic
from kafka.errors import KafkaError, TopicAlreadyExistsError
from cy_query.conf import settings
from cy_query.core import signing

__all__ = ['create_signed_token', 'verify_signed_token','Kafka']


def create_signed_token(data, max_age=None):
    """
    创建签名 Token
    :param data: 要签名的数据（字典）
    :param max_age: Token 有效时间（秒）
    :return: 签名后的字符串
    """
    return signing.dumps(data, key=settings.SECRET_KEY, salt='myapp-auth', compress=True)


# 验证并获取 Token 数据
def verify_signed_token(token, max_age=None):
    """
    验证签名 Token
    :param token: 要验证的 Token
    :param max_age: Token 有效时间（秒）
    :return: 原始数据或 None（如果验证失败）
    """
    try:
        return signing.loads(token, key=settings.SECRET_KEY, salt='myapp-auth', max_age=max_age)
    except (signing.BadSignature, signing.SignatureExpired):
        raise


class Kafka:
    def __init__(self, host, port=9092):
        self.host = host
        self.port = port
        self.bootstrap_servers = [host + ':' + str(port)]
        self._producer_conn = None
        self._admin_client = None
        self._consumer_conns = []

    def _get_admin_client(self) -> KafkaAdminClient:
        """获取 AdminClient（懒加载）"""
        if self._admin_client is None:
            self._admin_client = KafkaAdminClient(
                bootstrap_servers=self.bootstrap_servers
            )
        return self._admin_client

    def create_topic(self, topic_name, num_partitions=1, replication_factor=1):
        """
        创建主题
        :param topic_name: 主题名称
        :param num_partitions: 分区数（默认1）
        :param replication_factor: 副本数（默认1，生产环境建议≥2）
        :return:
        """
        try:
            admin_client = self._get_admin_client()
            topic_list = [
                NewTopic(name=topic_name, num_partitions=num_partitions, replication_factor=replication_factor)]
            admin_client.create_topics(new_topics=topic_list, validate_only=False)
            return True
        except TopicAlreadyExistsError:
            raise TopicAlreadyExistsError(f"主题 {topic_name} 已存在")
        except KafkaError as e:
            raise KafkaError(e)

    def list_topics(self):
        try:
            admin_client = self._get_admin_client()
            topics_metadata = admin_client.list_topics()
            return topics_metadata
        except KafkaError as e:
            raise

    @property
    def producer_conn(self):
        if self._producer_conn is None:
            self._producer_conn = KafkaProducer(
                bootstrap_servers=self.bootstrap_servers,
                value_serializer=lambda v: orjson.dumps(v),
                acks='all',
                key_serializer=lambda k: orjson.dumps(k) if k is not None else None,
                retries=3,  # 重试次数
            )
        return self._producer_conn

    def send(self, topic_name, message, key=None, callback: Optional[Callable] = None, flush=False,
             error_callback: Optional[Callable] = None):
        if not message:
            return
        message = create_signed_token(message)  # 对消息进行签名
        future = self.producer_conn.send(topic_name, value=message, key=key)
        try:
            if callback or error_callback:
                def on_completion(record_metadata_or_exception):
                    if isinstance(record_metadata_or_exception, Exception):
                        if error_callback:
                            error_callback(record_metadata_or_exception)
                        else:
                            raise RuntimeError(f"Failed to send message: {e}")
                    else:
                        if callback:
                            callback(record_metadata_or_exception)

                future.add_callback(on_completion)
            if flush:
                self.producer_conn.flush()
        except KafkaError as e:
            if error_callback:
                error_callback(e)
            else:
                raise RuntimeError(f"Failed to send message: {e}")

    def consumer(self, topic_name, group_id, auto_offset_reset='earliest'):
        """
        消费者
        :param topic_name: 主题名称
        :param group_id: 消费者组
        :param auto_offset_reset: 当没有初始偏移量或偏移量超出范围时，消费者会从指定的位置开始消费消息。设置为 earliest 会从最早的消息开始。
        :return:
        """
        consumer_conn = KafkaConsumer(
            topic_name,
            bootstrap_servers=self.bootstrap_servers,
            group_id=group_id,  # 建议使用新的group_id测试（避免旧偏移量影响）
            auto_offset_reset=auto_offset_reset,
            value_deserializer=lambda m: orjson.loads(m)
        )
        self._consumer_conns.append(consumer_conn)
        for message in consumer_conn:
            try:
                yield {'offset': message.offset, 'value': verify_signed_token(message.value), 'key': message.key,
                       'count': message.count, 'message': message}
            except (signing.BadSignature, signing.SignatureExpired):
                # 如果校验失败说明被修改了
                continue

    def __enter__(self):
        return self

    def __exit__(self, *args):
        self.close()

    def close(self) -> None:
        """关闭所有连接"""
        if self._producer_conn:
            try:
                self._producer_conn.close()
            except:
                pass
        if self._admin_client:
            try:
                self._admin_client.close()
            except:
                pass
        for consumer_conn in self._consumer_conns:
            try:
                consumer_conn.close()
            except:
                pass
